package org.example.nio;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class NonblockingServer extends Thread {

  private static final Logger logger = LoggerFactory.getLogger(NonblockingServer.class);

  private final Selector serverSocketSelector;

  private final List<Selector> channelSocketReadSelectors;

  private final AtomicInteger lbCount = new AtomicInteger(0);

  private final ServerSocketChannel serverSocketChannel;
  private final Executor executor;

  public NonblockingServer(int port) throws IOException {
    this(port, 0);
  }

  public NonblockingServer(int port, int worker) throws IOException {
    serverSocketSelector = Selector.open();

    if (worker <= 0) {
      worker = Runtime.getRuntime().availableProcessors() * 2;
    }

    channelSocketReadSelectors = new ArrayList<>(worker);

    for (int i = 0; i < worker; i++) {
      channelSocketReadSelectors.add(Selector.open());
    }

    serverSocketChannel = ServerSocketChannel.open();
    serverSocketChannel.configureBlocking(false);
    serverSocketChannel.bind(new InetSocketAddress(port));
    serverSocketChannel.register(serverSocketSelector, SelectionKey.OP_ACCEPT);

    executor = Executors.newFixedThreadPool(1 + channelSocketReadSelectors.size());
  }

  private Selector lbReadSelector() {
    int i = lbCount.incrementAndGet();
    return channelSocketReadSelectors.get(Math.abs(i) % channelSocketReadSelectors.size());
  }

  private void serverSocketSelector() {
    logger.info("ServerSocketSelector Start");

    try {
      while (serverSocketSelector.select() > 0) {
        logger.info("ServerSocketSelector Select");

        try {
          Iterator<SelectionKey> iterator = serverSocketSelector.selectedKeys().iterator();
          while (iterator.hasNext()) {
            SelectionKey selectedKey = iterator.next();
            SocketChannel socketChannel = ((ServerSocketChannel) selectedKey.channel()).accept();
            socketChannel.configureBlocking(false);
            socketChannel.register(lbReadSelector(), SelectionKey.OP_READ);
            if (logger.isDebugEnabled()) {
              logger.debug("socketChannel channelSocketReadSelector register OP_READ");
            }

            iterator.remove();
          }
        }
        catch (IOException e) {
          logger.error("ServerSocketSelector ERROR", e);
          throw new RuntimeException(e);
        }
      }
    }
    catch (Exception e) {
      logger.error("ServerSocketSelector ERROR", e);
      throw new RuntimeException(e);
    }
    logger.info("ServerSocketSelector 结束了");
  }

  private void clientSocketReadSelector(Selector channelSocketReadSelector) {
    logger.info("ClientSocketReadSelector Start");
    ByteBuffer readBuffer = ByteBuffer.allocate(1024);
    while (true) {
      try {
        if (channelSocketReadSelector.select(1) <= 0) {
          continue;
        }

        if (logger.isDebugEnabled()) {
          logger.debug("ClientSocketReadSelector Start READ");
        }

        Iterator<SelectionKey> iterator = channelSocketReadSelector.selectedKeys().iterator();
        while (iterator.hasNext()) {
          SelectionKey selectedKey = iterator.next();
          SocketChannel socketChannel = (SocketChannel) selectedKey.channel();
          if (!selectedKey.isValid() || !selectedKey.isReadable()) {
            iterator.remove();
            if (socketChannel.isOpen()) {
              socketChannel.close();
            }
            continue;
          }
          int num = socketChannel.read(readBuffer);
          if (num > 0) {
            readBuffer.flip();
            if (logger.isDebugEnabled()) {
              logger.debug("body READ \n{}", new String(readBuffer.array(), 0, num));
            }
            readBuffer.flip();
            readBuffer.clear();
          }
          else if (num == -1) {
            logger.info("ClientSocketReadSelector cancel");
            socketChannel.close();
            iterator.remove();
            continue;
          }

          byte[] results = """
              HTTP/1.1 200 OK\r
              Content-Type: text/html; charset=utf-8\r
              Connection: keep-alive\r
              \r
              <h1>哈哈哈哈</h1>\r<h2>%s</h2>
              """.formatted(UUID.randomUUID().toString()).getBytes();

          ByteBuffer byteBuffer = ByteBuffer.allocate(results.length);
          byteBuffer.put(results);
          byteBuffer.flip();

          socketChannel.write(byteBuffer);
          byteBuffer.clear();
          socketChannel.shutdownOutput();
          iterator.remove();
        }
      }
      catch (IOException e) {
        logger.error("ClientSocketReadSelector ERROR", e);
      }
    }
  }

  @Override
  public void run() {
    executor.execute(this::serverSocketSelector);
    for (Selector channelSocketReadSelector : channelSocketReadSelectors) {
      executor.execute(() -> clientSocketReadSelector(channelSocketReadSelector));
    }
  }
}
